#include <iostream>
#include <vector>
#include <cmath>
#include <fstream>
#include <cstdlib>
#include <Eigen/Dense>
#include <complex>
using namespace Eigen;
using namespace std;

const double PI = 3.14159265358979323846;

double lambda(double w, int n, int k) {
    return 1 - 2.0 * w * pow(sin(k * PI / (2 * n)), 2);
}

double c(int k, int n) {
    return pow(cos(k * PI / (2 * n)), 2);
}

double s(int k, int n) {
    return pow(sin(k * PI / (2 * n)), 2);
}

MatrixXd compute_matrix(double w, int n, int k, int v1, int v2) {
    int k_add = n - k;
    double c1 = pow(lambda(w, n, k), v1 + v2) * s(k, n);
    double c2 = pow(lambda(w, n, k_add), v1) * pow(lambda(w, n, k), v2) * s(k_add, n);
    double c3 = pow(lambda(w, n, k_add), v1) * pow(lambda(w, n, k), v2) * c(k, n);
    double c4 = pow(lambda(w, n, k), v1 + v2) * c(k_add, n);
    return (MatrixXd(2, 2) << c1, c2, c3, c4).finished();
}

double spectral_radius(const MatrixXd& matrix) {
    ComplexEigenSolver<MatrixXd> ces(matrix);
    VectorXcd eigenvalues = ces.eigenvalues();
    double max_abs = 0.0;
    for (int i = 0; i < eigenvalues.size(); ++i) {
        double abs_val = eigenvalues[i].real() * eigenvalues[i].real() + eigenvalues[i].imag() * eigenvalues[i].imag();
        if (abs_val > max_abs) {
            max_abs = abs_val;
        }
    }
    return sqrt(max_abs);
}

void save_to_file(const vector<int>& k_values, const vector<double>& c1, const vector<double>& c2, const vector<double>& c3, const vector<double>& c4, const vector<double>& rho, const string& filename) {
    ofstream file(filename);
    if (!file.is_open()) {
        cerr << "Error opening file " << filename << endl;
        exit(1);
    }
    for (size_t i = 0; i < k_values.size(); ++i) {
        file << k_values[i] << " " << c1[i] << " " << c2[i] << " " << c3[i] << " " << c4[i] << " " << rho[i] << endl;
    }
    file.close();
}

int main() {
    int n = 128;
    double w = 2.0 / 3.0;
    vector<pair<int, int>> args = {{0, 0}, {0, 2}, {1, 1}, {2, 0}, {2, 2}, {4, 0}};
    
    for (const auto& arg : args) {
        int v1 = arg.first;
        int v2 = arg.second;
        
        vector<int> k_values;
        vector<double> c1_values, c2_values, c3_values, c4_values, rho_values;
        
        for (int k = 0; k <= n; ++k) {
            MatrixXd matrix = compute_matrix(w, n, k, v1, v2);
            double rho = spectral_radius(matrix);
            
            k_values.push_back(k);
            c1_values.push_back(matrix(0, 0));
            c2_values.push_back(matrix(0, 1));
            c3_values.push_back(matrix(1, 0));
            c4_values.push_back(matrix(1, 1));
            rho_values.push_back(rho);
        }
        
        // Save data to file
        string filename = "data_v1_" + to_string(v1) + "_v2_" + to_string(v2) + ".dat";
        save_to_file(k_values, c1_values, c2_values, c3_values, c4_values, rho_values, filename);
        
        // Generate gnuplot script
        string gnuplot_script_name = "plot_v1_" + to_string(v1) + "_v2_" + to_string(v2) + ".gnuplot";
        ofstream gnuplot_script(gnuplot_script_name);
        gnuplot_script << "set terminal pngcairo enhanced font 'Verdana,10' size 800,600\n";
        string output_filename = "plot_v1_" + to_string(v1) + "_v2_" + to_string(v2) + ".png";
        gnuplot_script << "set output '" << output_filename << "'\n";
        gnuplot_script << "set xlabel 'Wavenumber k'\n";
        gnuplot_script << "set ylabel 'Damping Coefficients'\n";
        string title = "nu_1 = " + to_string(v1) + ", nu_2 = " + to_string(v2);
        gnuplot_script << "set title '" << title << "'\n";
        gnuplot_script << "plot '" << filename << "' using 1:2 title 'c1' with lines, '' using 1:3 title 'c2' with lines, '' using 1:4 title 'c3' with lines, '' using 1:5 title 'c4' with lines, '' using 1:6 title 'rho' with lines\n";
        gnuplot_script.close();
        
        // Execute gnuplot
        string command = "gnuplot " + gnuplot_script_name;
        system(command.c_str());
    }
    
    return 0;
}